#!/usr/bin/python3
# ==============================================================================
#
# Copyright (C) 2022 Sophgo Technologies Inc.  All rights reserved.
#
# TPU-MLIR is licensed under the 2-Clause BSD License except for the
# third-party components.
#
# ==============================================================================

import sys
import argparse as arg

from profile_helper.interface import *


def profileArgParser():
    parser = arg.ArgumentParser()
    # yapf: disable
    parser.add_argument("--mode", type=str, choices=["time", "command", "check", "sim"],
                        help="time: parse profile data to visualize. "
                        "sim: parse net_stat.sim in profile to visualize. "
                        "command: parse static commands in profile to analyse the command params. ",
                        default="time")
    parser.add_argument("--arch", type=str, choices=["BM1684", "BM1684X"], help="chip arch", default="BM1684")
    parser.add_argument("--format", type=str, choices=["html", "csv", "console", "layer"], help="output format", default="html")
    parser.add_argument("--test", type=str, help="only for mode=command, check some condition for command", default="")
    parser.add_argument("--option", type=str, help="profile extra options, format is 'key1=value1,key2=value2,...'",
                        default="")
    parser.add_argument("input_dir", type=str, help="profile_data dir generated by bmruntime")
    parser.add_argument("output_dir", type=str, help="output dir that contains parsing result", default="profile_out")
    # yapf: enable
    return parser


if __name__ == "__main__":
    parser = profileArgParser()
    args = parser.parse_args()
    if args.mode == "time":
        bmprofile_analyze(args.input_dir, args.output_dir, args.format, args.option)
    elif args.mode == "sim":
        bmprofile_simulate(args.input_dir, args.output_dir, args.format, args.option)
    elif args.mode == "command":
        bmprofile_parse_command(args.input_dir, args.output_dir, args.test, arch=args.arch)
    elif args.mode == "check":
        bmprofile_check_command(args.input_dir, args.output_dir, args.test, arch=args.arch)
    else:
        assert 0, "Not support mode={}".format(args.mode)
